import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

# torch.manual_seed(1)    # reproducible

x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)  # x data (tensor), shape=(100, 1)
y = x.pow(2) + 0.2 * torch.rand(x.size())               # noisy y data (tensor), shape=(100, 1)

# define the network
class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden) # hidden layer
        self.out = torch.nn.Linear(n_hidden, n_output)     # output layer

    def forward(self, x):
        x = torch.tanh(self.hidden(x)) # sigmoid/tanh/relu activation function for hidden layer
        x = self.out(x)            # linear output
        return x

net = Net(n_feature=1, n_hidden=10, n_output=1) # create a network
print(net)  # net architecture

optimizer = torch.optim.SGD(net.parameters(), lr=0.2)
loss_func = torch.nn.MSELoss()  # this is for regression mean squared loss

plt.ion()   # something about plotting

for t in range(1000):
    prediction = net(x)             # input x and predict based on x
    loss = loss_func(prediction, y) # must be (1. nn output, 2. target)
    optimizer.zero_grad()           # clear gradients for next train
    loss.backward()                 # backpropagation, compute gradients
    optimizer.step()                # apply gradients

    if t % 5 == 0:
        # plot and show learning process
        plt.cla()
        plt.scatter(x.data.numpy(), y.data.numpy())
        plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
        plt.text(0.5, 0, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color':  'red'})
        plt.pause(0.1)

print('training done.\n')

# save and restore method i:
torch.save(net.state_dict(), 'my_nn_net.params') # only save the parameters
net1 = Net(n_feature=1, n_hidden=10, n_output=1) # create a network with the same arch like the saved one.
net1.load_state_dict(torch.load('my_nn_net.params'))
prediction = net1(x)

# save and restore method ii:
torch.save(net, 'my_nn_net.net')   # save the whole net.
net2 = torch.load('my_nn_net.net') # restore and crete a new net like the saved one.
prediction = net2(x)

# show parameters.
print('show the parameters:')
params = net.state_dict() 
for k,v in params.items():
    print(k)                # show tensor key name
    print(v.numpy(), '\n')  # show tensor values

plt.ioff()
plt.show()
